#!/Pub/Users/wangyk/miniconda3/envs/pyscenic/bin/python3.10
import argparse

wrp = '''
    inferCNVpy命令行分析，利用infercnvpy进行CNV分析。在key_column列中，需要将上皮细胞分为多个亚群。比如说，一个seurat对象中包含上皮与成纤维细胞，上皮细胞分为0-9共计十个亚群，然后在本次分析中，会与对照细胞亚群进行inferCNVpy分析，识别0-9当中哪些上皮细胞亚群是恶性的。
    v2版本主要是改变了AnnData的生成方式，之前是利用seurat对象生成一个h5ad文件，偶尔报错，速度很慢，新版本直接提取表达counts到feather文件，细胞信息到tsv文件，利用两个文件与scanpy直接生成一个AnnData文件用于后续分析，大大加快运行速度。
    v3版本支持计算每个基因的cnvs得分。

命令参考：
inferCNV_py -s /Pub/Users/wangyk/project/Poroject/F240510002_OV_marco/out/SC_part/05.macro_sub_cellchat/ss_m_sub_anno_ready.rds -od Pub/Users/wangyk/project/Poroject/F240510002_OV_marco/tes2/ -k SubType -r TRM -w 12 -he 9

'''


parser = argparse.ArgumentParser(prog='inferCNV_py',
                                 description=wrp, 
                                 formatter_class=argparse.RawTextHelpFormatter)

parser.add_argument('-h5', '--h5ad', help="h5ad文件路径。scanpy读取的单细胞分析h5ad对象",
                    default='', metavar='')
parser.add_argument('-s', '--seurat_obj', help="seurat对象文件路径。如果没有现成的h5ad对象，则尝试利用seurat对象的rds或rdata文件进行转化，会在结果路径中生成一个表达谱feather文件与细胞注释tsv文件",
                    default=False, metavar='')
parser.add_argument('-od', '--output_dir', help="结果输出路径",
                    default='./', metavar='')
parser.add_argument('-k', '--key_column',
                    help='用于识别细胞分群的列名，必须包含在AnnData的obs中，或者说seurat中meta.data中的某一列。比如用cell_type这一列的注释信息对细胞进行分群，其中包含0、1、2、3、T细胞亚群，然后使用T细胞亚群作为reference，对0、1、2、3亚群的细胞进行CNV识别。',
                    default='SubType',
                    metavar='')
parser.add_argument('-r', '--reference', help="key_column列中，用于做对照的细胞类型。",
                    default='TRM', metavar='')
parser.add_argument('-w', '--width', help="数字，成图宽，默认为9",
                    default=9, metavar='', type=int)
parser.add_argument('-he', '--height', help="数字，成图高，默认为7",
                    default=7, metavar='', type=int)
parser.add_argument('-p', '--plot', help="逻辑值，对结果进行可视化，默认True",
                    default=True, metavar='')

args = parser.parse_args()

import scanpy as sc
import logging
import pyarrow.feather as feather
import infercnvpy as cnv
import sys
import os
import pandas as pd


def check_file_existence(file_path: str):
    """
    检查文件是否存在。

    Args:
        file_path (str): 要检查的文件路径。

    Raises:
        SystemExit: 如果文件不存在，记录错误信息并退出程序。
    """
    if not os.path.isfile(file_path):
        logger.error("未找到文件，请检查路径：%s", file_path)
        sys.exit(1)


def ensure_directory_exists(directory_path: str):
    """
    确保指定目录存在，如果不存在则创建。

    Args:
        directory_path (str): 要检查的目录路径。
    """
    if not os.path.exists(directory_path):
        os.makedirs(directory_path)

# 配置日志
def setup_logging():
    logging.basicConfig(level=logging.INFO,
                        format='%(asctime)s - %(levelname)s - %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S')
    return logging.getLogger("inferCNV_py")

# 尝试将seurat对象转换为feather文件
def convert_seurat_to_feather(seurat_obj_path: str, output_dir: str):
    r_file = "/Pub/Users/wangyk/Project_wangyk/Codelib_YK/codeLibrary_yk/py/SingleCell/utils4cmd/inferCNV_py/v3/get_exprs_and_cell_info.r"
    os.environ['R_HOME'] = "/Pub/Apps/Cellar/R/4.1.2/lib/R"
    try:
        os.system(f'{r_file} {seurat_obj_path} {output_dir} -t counts')
    except Exception as e:
        logger.error("转换seurat对象到feather文件时出错：%s", str(e))
        sys.exit(1)

# 从表达文件和细胞注释信息文件来构建一个AnnData对象
def get_anndata_from_feather(expr_feather: str, cell_info_tsv: str):
    """
    从表达文件和细胞注释信息文件来构建一个AnnData对象:
    
    Args:
        expr_feather (str): 表达谱feather文件。要求基因第一列样本在其他列；
        cell_info_tsv (str): 细胞注释信息tsv文件，第一列为细胞id，其余列为注释信息列。
    """
    # 1. 读取 counts 矩阵的 feather 文件
    counts_matrix = feather.read_feather(expr_feather)
    counts_matrix = counts_matrix.rename(columns={counts_matrix.columns[0]: 'gene'})
    counts_matrix.set_index('gene', inplace=True)
    # 2. 读取细胞注释文件
    cell_annotations = pd.read_csv(cell_info_tsv, sep='\t', index_col=0)
    # 3. 创建 AnnData 对象，确保数据为 cells x genes
    adata = sc.AnnData(X=counts_matrix.T)
    # 4. 将细胞注释信息加入 AnnData 对象
    adata.obs = cell_annotations
    return adata

# 进行cnv分析
def perform_cnv_analysis(adata, key_column, reference):
    logger.info("Adding genomic position information to AnnData Obj ...")
    cnv.io.genomic_position_from_gtf(
        "/Pub/Data/Data_Center/Annotation/gencode.v38.annotation.gtf", adata)

    # cnv分析
    cnv.tl.infercnv(
        adata,
        reference_key=key_column,
        reference_cat=reference,
        n_jobs=20
    )
    # 计算cnv score
    adata.obs['cell'] = adata.obs.index
    cnv.tl.cnv_score(adata, groupby='cell', key_added='cnv_score')
    return adata

# 保存结果
def save_results(adata, output_dir, key_column, reference, width, height, plot):
    os.chdir(output_dir)
    # save fig
    cnv.pl.chromosome_heatmap(adata, groupby=key_column, figsize=(
        width, height), save='Fig_inferCNVpy.pdf', dendrogram=True)
    # save score
    adata.obs['cell'] = adata.obs.index
    feather.write_feather(adata.obs, dest='data_frame_meta.data.feather')

    if plot:
        r_heat_plot_cmd = "/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/SC/inferCNV_py/inferCNV_py_plot_in_r.r"
        try:
            os.system(
                f"{r_heat_plot_cmd} -k {key_column} -r {reference} -od {output_dir}")
        except Exception as e:
            logger.error("绘制结果图时出错：%s", str(e))
            sys.exit(1)

# 主函数
def main():
    ensure_directory_exists(args.output_dir)
    global logger
    logger = setup_logging()

    if args.seurat_obj:
        convert_seurat_to_feather(args.seurat_obj, args.output_dir)

    if os.path.exists(args.h5ad):
        adata = sc.read_h5ad(args.h5ad)
    else:
        logger.info("生成表达矩阵feather，添加cellinfo ...")
        adata = get_anndata_from_feather(
            os.path.join(args.output_dir, "seurat_obj_expression.feather"),
            os.path.join(args.output_dir, "seurat_obj_cellInfo.tsv")
        )

    if args.key_column not in adata.obs.columns.values:
        logger.error("指定列在AnnData.obs中不存在，请检查")
        sys.exit(1)

    adata = perform_cnv_analysis(adata, args.key_column, args.reference)
    save_results(adata, args.output_dir, args.key_column, args.reference, args.width, args.height, args.plot)

    logger.info("分析完成")

if __name__ == "__main__":
    main()